#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# limit the number of cpus used by high performance libraries

from cgi import test
from tkinter.font import names
import rospy
import roslib

from std_msgs.msg import Header
from sensor_msgs.msg import Image
from yolov5_ros_msgs.msg import BoundingBox, BoundingBoxes

import os
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"

import sys
sys.path.insert(0, './yolov5')
sys.path.append("..")
from pathlib import Path

sys.path.append(str(Path(__file__, "..", "yolov5").resolve()))


import argparse
import os
import platform
import shutil
import time
from pathlib import Path
import numpy as np
import cv2
import torch
import torch.backends.cudnn as cudnn

from yolov5.models.experimental import attempt_load
from yolov5.utils.downloads import attempt_download
from yolov5.models.common import DetectMultiBackend
from yolov5.utils.datasets import LoadImages, LoadStreams, VID_FORMATS
from yolov5.utils.general import (LOGGER, check_img_size, imshow, non_max_suppression, print_args, scale_coords,
                                  check_imshow, xyxy2xywh, increment_path, strip_optimizer, colorstr)
from yolov5.utils.torch_utils import select_device, time_sync
from yolov5.utils.plots import Annotator, colors, save_one_box
from deep_sort.utils.parser import get_config
from deep_sort.deep_sort import DeepSort

import matplotlib.pyplot as plt
from PIL import Image as pilimage
from cv_bridge import CvBridge 

import albumentations as albu
import torchvision.transforms as transforms

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # yolov5 deepsort root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative



# new add 

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    # Resize and pad image while meeting stride-multiple constraints
    # print(im.shape)
 
    # print(im)
    # print("im.shape letterbox")
    shape = im.shape[:2]  # current shape [height, width]
    # print(shape)
    # print("shape is in leetr")
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
        print("isin")
    # print(new_shape)
    # print("new_shape")
    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    # print(r)
    # print("r")
    if not scaleup:  # only scale down, do not scale up (for better val mAP)
        r = min(r, 1.0)

    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    # print(new_unpad)
    # print("is new_uppad")
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    dw /= 2  # divide padding into 2 sides
    dh /= 2
    # print(shape[::-1])
    # print("::-1")
    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
        print("here_t")
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    # print(im.shape)
    # print("today")
    return im, ratio, (dw, dh)





class Yolo_Deepsort_Dect:
    def __init__(self):

        # load parameters
        image_topic = rospy.get_param(
            '~image_topic', '/camera/color/image_raw')
        pub_topic = rospy.get_param('~pub_topic', '/yolov5/BoundingBoxes')
        self.camera_frame = rospy.get_param('~camera_frame', '')

        parser = argparse.ArgumentParser()
        parser.add_argument('--yolo_model', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
        parser.add_argument('--deep_sort_model', type=str, default='osnet_x0_5')
        parser.add_argument('--source', type=str, default='0', help='source')  # file/folder, 0 for webcam
        parser.add_argument('--output', type=str, default='inference/output', help='output folder')  # output folder
        parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
        parser.add_argument('--conf-thres', type=float, default=0.5, help='object confidence threshold')
        parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
        parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
        parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
        parser.add_argument('--show-vid', action='store_true', help='display tracking video results')
        parser.add_argument('--save-vid', action='store_true', help='save video tracking results')
        parser.add_argument('--save-txt', action='store_true', help='save MOT compliant results to *.txt')
        # class 0 is person, 1 is bycicle, 2 is car... 79 is oven
        parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 16 17')
        parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
        parser.add_argument('--augment', action='store_true', help='augmented inference')
        parser.add_argument('--update', action='store_true', help='update all models')
        parser.add_argument('--evaluate', action='store_true', help='augmented inference')
        parser.add_argument("--config_deepsort", type=str, default="/home/wangbingtao/new_ws/src/yolo_deepsort_ros/scripts/Yolov5_DeepSort_OSNet/deep_sort/configs/deep_sort.yaml")
        parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
        parser.add_argument('--visualize', action='store_true', help='visualize features')
        parser.add_argument('--max-det', type=int, default=1000, help='maximum detection per image')
        parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
        parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
        parser.add_argument('--project', default=ROOT / 'runs/track', help='save results to project/name')
        parser.add_argument('--name', default='exp', help='save results to project/name')
        parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
        # print("here1")
        opt, unparsed = parser.parse_known_args()
        # opt = parser.parse_args()
        # print("here2")
        opt.show_vid = True
        opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
        self.getImageStatus = False

        self.opt = opt
        # detect 参数引入 from opt 
        out, source, yolo_model, deep_sort_model, show_vid, save_vid, save_txt, imgsz, evaluate, half, \
            project, exist_ok, update, save_crop = \
            opt.output, opt.source, opt.yolo_model, opt.deep_sort_model, opt.show_vid, opt.save_vid, \
            opt.save_txt, opt.imgsz, opt.evaluate, opt.half, opt.project, opt.exist_ok, opt.update, opt.save_crop
        webcam = source == '0' or source.startswith(
            'rtsp') or source.startswith('http') or source.endswith('.txt')
        # opt OK 

        self.nums = 0
        # Initialize
        self.device = select_device(opt.device)
        device = self.device

        self.half = half
        half &= device.type != 'cpu'  # half precision only supported on CUDA

        # The MOT16 evaluation runs multiple inference streams in parallel, each one writing to
        # its own .txt file. Hence, in that case, the output folder is not restored
        if not evaluate:
            if os.path.exists(out):
                pass
                shutil.rmtree(out)  # delete output folder
            os.makedirs(out)  # make new output folder

        # Directories
        if type(yolo_model) is str:  # single yolo model
            exp_name = yolo_model.split(".")[0]
        elif type(yolo_model) is list and len(yolo_model) == 1:  # single models after --yolo_model
            exp_name = yolo_model[0].split(".")[0]
        else:  # multiple models after --yolo_model
            exp_name = "ensemble"
        exp_name = exp_name + "_" + deep_sort_model.split('/')[-1].split('.')[0]
        save_dir = increment_path(Path(project) / exp_name, exist_ok=exist_ok)  # increment run if project name exists
        (save_dir / 'tracks' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

        # Load model
        self.model = DetectMultiBackend(yolo_model, device=device, dnn=opt.dnn)
        stride, names, pt = self.model.stride, self.model.names, self.model.pt
        imgsz = check_img_size(imgsz, s=stride)  # check image size

        self.stride = stride
        self.pt = pt
        # Half
        half &= pt and device.type != 'cpu'  # half precision only supported by PyTorch on CUDA
        if pt:
            self.model.model.half() if half else self.model.model.float()

        # Set Dataloader
        vid_path, vid_writer = None, None
        # Check if environment supports image displays
        if show_vid:
            show_vid = check_imshow()
            self.show_vid = show_vid
        # print("here333")

        # # Dataloader
        # if webcam:
        #     show_vid = check_imshow()
        #     cudnn.benchmark = True  # set True to speed up constant image size inference
        #     dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
        #     nr_sources = len(dataset)
        # else:
        #     dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
        #     nr_sources = 1
        # vid_path, vid_writer, txt_path = [None] * nr_sources, [None] * nr_sources, [None] * nr_sources
        
        # 把Dataloader中的用到的变量拿出来
        nr_sources = 1

        # initialize deepsort
        cfg = get_config()
        cfg.merge_from_file(opt.config_deepsort)

        # Create as many trackers as there are video sources
        deepsort_list = []
        for i in range(nr_sources):
            deepsort_list.append(
                DeepSort(
                    deep_sort_model,
                    device,
                    max_dist=cfg.DEEPSORT.MAX_DIST,
                    max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
                    max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
                )
            )
        outputs = [None] * nr_sources
        #全局化
        self.outputs = outputs
        self.deepsor_list = deepsort_list

        # Get names and colors
        names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
        self.names = names
        # Run tracking
        self.model.warmup(imgsz=(1 if pt else nr_sources, 3, *imgsz))  # warmup
        # dt, seen = [0.0, 0.0, 0.0, 0.0], 0

        # 来个显示图像的容器 
        # self.im_reslut = None
        self.im_reslut = np.zeros((1080,1920,3),np.uint8)
        # self.im_reslut.fill(255)
        # self.im_reslut = np.array()


        # image subscribe
        self.color_sub = rospy.Subscriber(image_topic, Image, self.image_callback,
                                          queue_size=1, buff_size=52428800)
        # print("???")

        # 显示结果 

        # cv2.imshow("hi",self.im_reslut)
        # cv2.waitKey(1)

        # print("22222222222222222222")
        # print(self.im_reslut.shape)

        # output publishers
        self.position_pub = rospy.Publisher(
            pub_topic,  BoundingBoxes, queue_size=1)

        self.image_pub = rospy.Publisher(
            '/yolov5/detection_image',  Image, queue_size=1)

        # if no image messages
        while (not self.getImageStatus) :
            rospy.loginfo("waiting for image.")
            rospy.sleep(2)
    

    def image_callback(self, image):
        self.boundingBoxes = BoundingBoxes()
        self.boundingBoxes.header = image.header
        self.boundingBoxes.image_header = image.header
        self.getImageStatus = True
        self.color_image = np.frombuffer(image.data, dtype=np.uint8).reshape(
            image.height, image.width, -1)
        # self.color_image = cv2.cvtColor(self.color_image, cv2.COLOR_BGR2RGB)
        print(self.color_image.shape)
        print("well")
        # trans = transforms.ToTensor()
        # self.color_image = trans(self.color_image)
        # print(self.color_image.shape)

        # 判断次数 
        self.nums = self.nums + 1
        if self.nums > 1000:
            self.nums = 100

        dt, seen = [0.0, 0.0, 0.0, 0.0], 0

        img = self.color_image
        # img = im.swapaxes(1,2).swapaxes(0,1)
        img_org = img
        t1 = time_sync()

        # Padded resize
        img = letterbox(img, self.opt.imgsz, stride=self.stride, auto=self.pt)[0]

        # Convert
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        img = np.ascontiguousarray(img)

        im = img
        im = torch.from_numpy(im).to(self.device)
        im = im.half() if self.half else im.float()  # uint8 to fp16/32
        im /= 255.0  # 0 - 255 to 0.0 - 1.0
        if len(im.shape) == 3:
            im = im[None]  # expand for batch dim
        
        print(im.shape)
        t2 = time_sync()
        dt[0] += t2 - t1

        pred = self.model(im)
        t3 = time_sync()
        dt[1] += t3 - t2

        # xmin    ymin    xmax   ymax  confidence  class    name
        pred = non_max_suppression(pred, self.opt.conf_thres, self.opt.iou_thres, self.opt.classes, self.opt.agnostic_nms, max_det=self.opt.max_det)
        # print(pred)
        # print("OOO")
        dt[2] += time_sync() - t3


        ## deepsort start 
        names = self.names
        outputs = self.outputs
        deepsort_list = self.deepsor_list

        # Process detections
        for i, det in enumerate(pred):  # detections per image
            seen += 1
            # if webcam:  # nr_sources >= 1
            #     p, im0, _ = path[i], im0s[i].copy(), dataset.count
            #     p = Path(p)  # to Path
            #     s += f'{i}: '
            #     txt_file_name = p.name
            #     save_path = str(save_dir / p.name)  # im.jpg, vid.mp4, ...
            # else:
            #     p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
            #     p = Path(p)  # to Path
            #     # video file
            #     if source.endswith(VID_FORMATS):
            #         txt_file_name = p.stem
            #         save_path = str(save_dir / p.name)  # im.jpg, vid.mp4, ...
            #     # folder with imgs
            #     else:
            #         txt_file_name = p.parent.name  # get folder name containing current img
            #         save_path = str(save_dir / p.parent.name)  # im.jpg, vid.mp4, ...

            # txt_path = str(save_dir / 'tracks' / txt_file_name)  # im.txt
            # s += '%gx%g ' % im.shape[2:]  # print string
            # imc = im0.copy() if save_crop else im0  # for save_crop
            

            # 弄一个原图像
            im0 = img_org.copy()
            annotator = Annotator(im0, line_width=2, pil=not ascii)
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

                # # Print results
                # for c in det[:, -1].unique():
                #     n = (det[:, -1] == c).sum()  # detections per class
                #     s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                xywhs = xyxy2xywh(det[:, 0:4])
                confs = det[:, 4]
                clss = det[:, 5]

                # pass detections to deepsort
                t4 = time_sync()
                outputs[i] = deepsort_list[i].update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
                t5 = time_sync()
                dt[3] += t5 - t4
                print(outputs)
                print("outputs")
                # draw boxes for visualization
                if len(outputs[i]) > 0:
                    for j, (output) in enumerate(outputs[i]):   
                        # 这里outpust[i] 指的是一个数据流中的第i张图片，在image_call 中就一张图片，所以i=0，在后续decshow函数中也是
                        bboxes = output[0:4]
                        id = output[4]
                        cls = output[5]
                        conf = output[6]
                        # print(bboxes)
                        # print("bboxes")
                        # if save_txt:
                        #     # to MOT format
                        #     bbox_left = output[0]
                        #     bbox_top = output[1]
                        #     bbox_w = output[2] - output[0]
                        #     bbox_h = output[3] - output[1]
                        #     # Write MOT compliant results to file
                        #     with open(txt_path + '.txt', 'a') as f:
                        #         f.write(('%g ' * 10 + '\n') % (frame_idx + 1, id, bbox_left,  # MOT format
                        #                                        bbox_top, bbox_w, bbox_h, -1, -1, -1, i))

                        # if save_vid or save_crop or show_vid:  # Add bbox to image
                        #     c = int(cls)  # integer class
                        #     label = f'{id:0.0f} {names[c]} {conf:.2f}'
                        #     annotator.box_label(bboxes, label, color=colors(c, True))
                        #     if save_crop:
                        #         txt_file_name = txt_file_name if (isinstance(path, list) and len(path) > 1) else ''
                        #         save_one_box(bboxes, imc, file=save_dir / 'crops' / txt_file_name / names[c] / f'{id}' / f'{p.stem}.jpg', BGR=True)

                        if self.show_vid:  # Add bbox to image
                            c = int(cls)  # integer class
                            label = f'{id:0.0f} {names[c]} {conf:.2f}'
                            annotator.box_label(bboxes, label, color=colors(c, True))
                          


                # LOGGER.info(f'{s}Done. YOLO:({t3 - t2:.3f}s), DeepSort:({t5 - t4:.3f}s)')
                LOGGER.info(f'Done. YOLO:({t3 - t2:.3f}s), DeepSort:({t5 - t4:.3f}s)')

            else:
                deepsort_list[i].increment_ages()
                LOGGER.info('No detections')

            # Stream results
            im0 = annotator.result()
            self.im_reslut = im0.copy()
            img_rel = im0.copy()

            # self.test()
            # print("1111111111111111111111")
            # print(self.im_reslut.shape)
            # print("1111111111111111111111")
  
            # # cv2.imshow("yeah",im0)
            # # cv2.waitKey(0)  # 1 millisecond
            # show_vid = self.show_vid
            # # print(show_vid)
            # # print("show_vid")
            # if show_vid:
            #     # cv2.imshow(str(p), im0)
            #     # cv2.imshow("yeah",im0)
            #     # print(im0.shape)
            #     # plt.imshow(img_rel)
            #     # plt.show()
            #     # vel = pilimage.open(img_rel)
            #     # vel.show()
            #     # vel = pilimage.fromarray(img_rel)
            #     # vel.show()
            #     print("end sub")


                # print("im0.shape")
                # cv2.waitKey(1)  # 1 millisecond

            # # Save results (image with detections)
            # if save_vid:
            #     if vid_path[i] != save_path:  # new video
            #         vid_path[i] = save_path
            #         if isinstance(vid_writer[i], cv2.VideoWriter):
            #             vid_writer[i].release()  # release previous video writer
            #         if vid_cap:  # video
            #             fps = vid_cap.get(cv2.CAP_PROP_FPS)
            #             w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            #             h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            #         else:  # stream
            #             fps, w, h = 30, im0.shape[1], im0.shape[0]
            #         save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
            #         vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
            #     vid_writer[i].write(im0)

 

    # # Print results
    # t = tuple(x / seen * 1E3 for x in dt)  # speeds per image
    # LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS, %.1fms deep sort update \
    #     per image at shape {(1, 3, *imgsz)}' % t)
    # if save_txt or save_vid:
    #     s = f"\n{len(list(save_dir.glob('tracks/*.txt')))} tracks saved to {save_dir / 'tracks'}" if save_txt else ''
    #     LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
    # if update:
    #     strip_optimizer(yolo_model)  # update model (to fix SourceChangeWarning)


    #     #copy end
    #     boxs = pred.pandas().xyxy[0].values
        # if self.nums > 10 :
        #     self.dectshow(self.color_image, outputs, image.height, image.width)


        self.dectshow(img_rel, outputs, image.height, image.width)


    #     cv2.waitKey(3)

    # def test(self):
    #     print(self.outputs)
    #     print("yoyo test")

    def dectshow(self, img_rel, boxs, height, width):
        # img = org_img.copy()

        # if len(outputs[i]) > 0:
        #     for j, (output) in enumerate(outputs[i]):

        #         bboxes = output[0:4]
        #         id = output[4]
        #         cls = output[5]
        #         conf = output[6]

        # count = 0
        # for i in boxs:
        #     count += 1

        for box in boxs[0]:
            # print(box)
            # print("here box")
            boundingBox = BoundingBox()
            # boundingBox.probability = np.float64(box[6])
            # boundingBox.xmin = np.int64(box[0])
            # boundingBox.ymin = np.int64(box[1])
            # boundingBox.xmax = np.int64(box[2])
            # boundingBox.ymax = np.int64(box[3])
            # boundingBox.num = np.int16(box[4])
            # boundingBox.Class = box[5]

            boundingBox.probability = float(box[6])
            boundingBox.xmin = int(box[0])
            boundingBox.ymin = int(box[1])
            boundingBox.xmax = int(box[2])
            boundingBox.ymax = int(box[3])
            boundingBox.num = int(box[4])
            boundingBox.Class = str(box[5])
            print(boundingBox)
            print("boundbox")

            # print(boundingBox.xmin)
            # print(type(boundingBox.xmin))

            # print("boundingBox.xmin tete")
            # if box[-1] in self.classes_colors.keys():
            #     color = self.classes_colors[box[-1]]
            # else:
            #     color = np.random.randint(0, 183, 3)
            #     self.classes_colors[box[-1]] = color

            # cv2.rectangle(img, (int(box[0]), int(box[1])),
            #               (int(box[2]), int(box[3])), (int(color[0]),int(color[1]), int(color[2])), 2)
            # cv2.putText(img, box[-1],
            #             (int(box[0]), int(box[1])-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2, cv2.LINE_AA)

            
            self.boundingBoxes.bounding_boxes.append(boundingBox)
            # print(self.boundingBoxes)
            # print("final")
        self.position_pub.publish(self.boundingBoxes)
        self.publish_image(img_rel, height, width)
        # cv2.imshow('YOLOv5', img)

    def publish_image(self, imgdata, height, width):
        # imgdata = CvBridge.cv2_to_imgmsg(imgdata,"bgr8")
        # self.color_image = cv2.cvtColor(self.color_image, cv2.COLOR_BGR2RGB)
        # imgdata = cv2.cvtColor(imgdata, cv2.COLOR_RGB2BGR)
        # 如果 cvtColor没有转换rgb格式为bgr格式 img_msg编码格式为rbg8，如果转换了 即用rbg8
        # 目前 imgdata格式为 rbg 一般来说 img_msg为bgr格式 记得带8 bgr8 

        image_temp = Image()
        header = Header(stamp=rospy.Time.now())
        header.frame_id = self.camera_frame
        image_temp.height = height
        image_temp.width = width
        image_temp.encoding = 'bgr8'
        image_temp.data = np.array(imgdata).tobytes()
        image_temp.header = header
        image_temp.step = width * 3
        self.image_pub.publish(image_temp)


def main():
    rospy.init_node('yolov5_deepsort_ros', anonymous=True)
    print("start")
    with torch.no_grad():
        yolo_dect = Yolo_Deepsort_Dect()
    rospy.spin()


if __name__ == "__main__":
   
    main()

















